import os
from pathlib import Path

import torch
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm

from config import config
from data.dataset import get_data_loaders
from models.cnn import CIFAR10CNN
from utils.logger import setup_logger

current_file = Path(__file__).name
logger = setup_logger(current_file)


def save_checkpoint(model, optimizer, epoch, acc, path):
    """保存检查点"""
    torch.save({
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'accuracy': acc,
    }, path)
    logger.info(f"Checkpoint saved to {path}")


def train_epoch(model, train_loader, criterion, optimizer, epoch, device):
    """单个训练epoch"""
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0

    pbar = tqdm(train_loader, desc=f"Train Epoch [{epoch + 1} / {config.EPOCHS}]", leave=False)
    for inputs, labels in pbar:
        inputs, labels = inputs.to(device), labels.to(device)

        # 前向传播
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)

        # 反向传播
        loss.backward()
        optimizer.step()

        # 统计指标
        running_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()

        # 更新进度条
        pbar.set_postfix({
            'Train Loss': f"{running_loss / (pbar.n + 1):.2f}",
            'Train Acc': f"{100. * correct / total:.2f}%"
        })

    epoch_loss = running_loss / len(train_loader)
    epoch_acc = 100. * correct / total
    return epoch_loss, epoch_acc


def test_epoch(model, test_loader, criterion, epoch, device):
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0

    with torch.no_grad():
        pbar = tqdm(test_loader, desc=f"Test Epoch [{epoch + 1} / {config.EPOCHS}]", leave=False)
        for images, labels in pbar:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)

            running_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

            # 更新进度条
            pbar.set_postfix({
                'Test Loss': f"{running_loss / (pbar.n + 1):.2f}",
                'Test Acc': f"{100. * correct / total:.2f}%"
            })

    test_loss = running_loss / len(test_loader)
    test_acc = 100 * correct / total

    return test_loss, test_acc


def train_model():
    # 创建模型保存目录
    os.makedirs(config.MODEL_SAVE_DIR, exist_ok=True)

    # 获取数据加载器
    train_loader, test_loader, classes = get_data_loaders()

    # 初始化模型、损失函数和优化器
    model = CIFAR10CNN(num_classes=len(classes)).to(config.DEVICE)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=config.LEARNING_RATE)

    best_accuracy = 0.0

    # 训练循环
    for epoch in range(config.EPOCHS):
        # 训练阶段
        train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer, epoch, config.DEVICE)

        # 测试阶段
        test_loss, test_acc = test_epoch(model, test_loader, criterion, epoch, config.DEVICE)

        # 保存最佳模型
        if test_acc > best_accuracy:
            best_accuracy = test_acc
            save_checkpoint(model, optimizer, epoch, best_accuracy, config.BEST_MODEL_PATH)
            logger.info(f"New best model saved with accuracy: {best_accuracy:.2f}%")

        logger.info(
            f"Epoch {epoch + 1}/{config.EPOCHS} - "
            f"Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}% - "
            f"Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.2f}%"
        )

        # 定期保存检查点
        if (epoch + 1) % config.SAVE_INTERVAL == 0:
            save_checkpoint(
                model, optimizer, epoch, test_acc,
                os.path.join(config.MODEL_SAVE_DIR, f"checkpoint_epoch{epoch}.pth")
            )

    logger.info(f"Training completed. Best test accuracy: {best_accuracy:.2f}%")


if __name__ == '__main__':
    train_model()
